#!/usr/bin/env python3
"""
MySQL数据库管理脚本
用于初始化、备份、恢复MySQL数据库
"""

import os
import sys
import subprocess
import pymysql
import logging
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from config.mysql_config import MySQLConfig

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class MySQLManager:
    """MySQL数据库管理器"""
    
    def __init__(self):
        self.config = MySQLConfig()
        self.config.validate_config()
    
    def test_connection(self):
        """测试MySQL连接"""
        try:
            connection = pymysql.connect(
                host=self.config.MYSQL_HOST,
                port=self.config.MYSQL_PORT,
                user=self.config.MYSQL_USER,
                password=self.config.MYSQL_PASSWORD,
                charset=self.config.MYSQL_CHARSET
            )
            
            with connection.cursor() as cursor:
                cursor.execute("SELECT VERSION()")
                version = cursor.fetchone()
                logger.info(f"MySQL连接成功，版本: {version[0]}")
            
            connection.close()
            return True
            
        except Exception as e:
            logger.error(f"MySQL连接失败: {e}")
            return False
    
    def create_database(self):
        """创建数据库"""
        try:
            connection = pymysql.connect(
                host=self.config.MYSQL_HOST,
                port=self.config.MYSQL_PORT,
                user=self.config.MYSQL_USER,
                password=self.config.MYSQL_PASSWORD,
                charset=self.config.MYSQL_CHARSET
            )
            
            with connection.cursor() as cursor:
                # 创建数据库
                cursor.execute(f"""
                    CREATE DATABASE IF NOT EXISTS {self.config.MYSQL_DATABASE} 
                    CHARACTER SET utf8mb4 
                    COLLATE utf8mb4_unicode_ci
                """)
                logger.info(f"数据库 {self.config.MYSQL_DATABASE} 创建成功")
            
            connection.commit()
            connection.close()
            return True
            
        except Exception as e:
            logger.error(f"创建数据库失败: {e}")
            return False
    
    def init_database(self):
        """初始化数据库（执行SQL脚本）"""
        sql_file = project_root / "scripts" / "init_mysql.sql"
        
        if not sql_file.exists():
            logger.error(f"初始化脚本不存在: {sql_file}")
            return False
        
        try:
            # 使用mysql命令行工具执行SQL脚本
            cmd = [
                "mysql",
                f"-h{self.config.MYSQL_HOST}",
                f"-P{self.config.MYSQL_PORT}",
                f"-u{self.config.MYSQL_USER}",
                f"-p{self.config.MYSQL_PASSWORD}",
                "--default-character-set=utf8mb4"
            ]
            
            with open(sql_file, 'r', encoding='utf-8') as f:
                result = subprocess.run(
                    cmd,
                    input=f.read(),
                    text=True,
                    capture_output=True
                )
            
            if result.returncode == 0:
                logger.info("数据库初始化成功")
                if result.stdout:
                    logger.info(f"输出: {result.stdout}")
                return True
            else:
                logger.error(f"数据库初始化失败: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"执行初始化脚本失败: {e}")
            return False
    
    def backup_database(self, backup_dir=None):
        """备份数据库"""
        if backup_dir is None:
            backup_dir = project_root / "backups"
        
        backup_dir = Path(backup_dir)
        backup_dir.mkdir(exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_file = backup_dir / f"{self.config.MYSQL_DATABASE}_backup_{timestamp}.sql"
        
        try:
            cmd = [
                "mysqldump",
                f"-h{self.config.MYSQL_HOST}",
                f"-P{self.config.MYSQL_PORT}",
                f"-u{self.config.MYSQL_USER}",
                f"-p{self.config.MYSQL_PASSWORD}",
                "--default-character-set=utf8mb4",
                "--single-transaction",
                "--routines",
                "--triggers",
                self.config.MYSQL_DATABASE
            ]
            
            with open(backup_file, 'w', encoding='utf-8') as f:
                result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, text=True)
            
            if result.returncode == 0:
                logger.info(f"数据库备份成功: {backup_file}")
                return str(backup_file)
            else:
                logger.error(f"数据库备份失败: {result.stderr}")
                backup_file.unlink(missing_ok=True)  # 删除失败的备份文件
                return None
                
        except Exception as e:
            logger.error(f"备份数据库失败: {e}")
            return None
    
    def restore_database(self, backup_file):
        """恢复数据库"""
        backup_file = Path(backup_file)
        
        if not backup_file.exists():
            logger.error(f"备份文件不存在: {backup_file}")
            return False
        
        try:
            cmd = [
                "mysql",
                f"-h{self.config.MYSQL_HOST}",
                f"-P{self.config.MYSQL_PORT}",
                f"-u{self.config.MYSQL_USER}",
                f"-p{self.config.MYSQL_PASSWORD}",
                "--default-character-set=utf8mb4",
                self.config.MYSQL_DATABASE
            ]
            
            with open(backup_file, 'r', encoding='utf-8') as f:
                result = subprocess.run(
                    cmd,
                    input=f.read(),
                    text=True,
                    capture_output=True
                )
            
            if result.returncode == 0:
                logger.info(f"数据库恢复成功: {backup_file}")
                return True
            else:
                logger.error(f"数据库恢复失败: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"恢复数据库失败: {e}")
            return False
    
    def show_status(self):
        """显示数据库状态"""
        try:
            connection = pymysql.connect(
                host=self.config.MYSQL_HOST,
                port=self.config.MYSQL_PORT,
                user=self.config.MYSQL_USER,
                password=self.config.MYSQL_PASSWORD,
                database=self.config.MYSQL_DATABASE,
                charset=self.config.MYSQL_CHARSET
            )
            
            with connection.cursor() as cursor:
                # 显示数据库信息
                cursor.execute("SELECT DATABASE(), VERSION(), NOW()")
                db_info = cursor.fetchone()
                logger.info(f"当前数据库: {db_info[0]}")
                logger.info(f"MySQL版本: {db_info[1]}")
                logger.info(f"当前时间: {db_info[2]}")
                
                # 显示表信息
                cursor.execute("SHOW TABLES")
                tables = cursor.fetchall()
                logger.info(f"数据库表数量: {len(tables)}")
                
                for table in tables:
                    cursor.execute(f"SELECT COUNT(*) FROM {table[0]}")
                    count = cursor.fetchone()[0]
                    logger.info(f"  - {table[0]}: {count} 条记录")
            
            connection.close()
            return True
            
        except Exception as e:
            logger.error(f"获取数据库状态失败: {e}")
            return False

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='MySQL数据库管理工具')
    parser.add_argument('action', choices=['test', 'init', 'backup', 'restore', 'status'],
                       help='要执行的操作')
    parser.add_argument('--backup-file', help='备份文件路径（用于恢复操作）')
    parser.add_argument('--backup-dir', help='备份目录路径')
    
    args = parser.parse_args()
    
    manager = MySQLManager()
    
    if args.action == 'test':
        success = manager.test_connection()
        sys.exit(0 if success else 1)
    
    elif args.action == 'init':
        if manager.test_connection():
            success = manager.create_database() and manager.init_database()
            sys.exit(0 if success else 1)
        else:
            sys.exit(1)
    
    elif args.action == 'backup':
        backup_file = manager.backup_database(args.backup_dir)
        sys.exit(0 if backup_file else 1)
    
    elif args.action == 'restore':
        if not args.backup_file:
            logger.error("恢复操作需要指定备份文件")
            sys.exit(1)
        success = manager.restore_database(args.backup_file)
        sys.exit(0 if success else 1)
    
    elif args.action == 'status':
        success = manager.show_status()
        sys.exit(0 if success else 1)

if __name__ == '__main__':
    main()